from .base_controller import BaseController
# from util.visualizer import Visualizer
import time
from data import create_dataset
from models import create_model,get_display_init
from visualizer.OpencvVis import CVis
import numpy as np
import copy
import models
class DisplayController(BaseController):
    def modify_commandline_options(parser,flagdict=None):
        parser.set_defaults(usemodel=False)
        parser.add_argument('--xreverse', type=int, default=1, help='')
        parser.add_argument('--yreverse', type=int, default=1, help='')
        parser.add_argument('--zreverse', type=int, default=1, help='')
        parser.add_argument('--ballradius', type=int, default=2, help='render ball radius')
        parser.add_argument('--cvh', type=int, default=128*6, help='heigh of display windows')
        parser.add_argument('--cvw', type=int, default=128*9, help='width of display windows')
        parser.add_argument('--cvz', type=int, default=300*3, help='deep of display windows')
        parser.add_argument('--showrot', type=bool, default=False, help='show rotation or not')
        parser.add_argument('--magnifyBlue', type=int, default=0, help='magnifyBlue')
        parser.add_argument('--freezerot', type=bool, default=False, help='freeze rotation or not')
        parser.add_argument('--normalizecolor', type=bool, default=True, help='use normalizecolor or not')
        parser.add_argument('--waittime', type=int, default=0, help='time of windows delay')
        parser.add_argument('--display_3d', type=str, default='yes', help='')
        parser.add_argument('--phase', type=str, default='display_zed_camera', help='')
        parser.add_argument('--pclnpoints', type=int, default=2000000, help='')
        parser.add_argument('--display_models', type=str, default='', help="The list of display moedels(split by ',')")
        parser.add_argument('--display_models_args', type=str, default='', help="The list of display moedels args(split by ',' , Arg format is ModelName:ArgName:ArgValue:ArgType)")

        return parser
    
    def __init__(self,opt):
        BaseController.__init__(self,opt)
        
        self.opt = opt   # get training options
        datasets = opt.dataset_mode
        self.datasets = create_dataset(opt)

        if self.opt.display_models == '':
            self.opt.usemodel = False
        else:
            self.models = []
            self._models = self.opt.display_models.split(',')
            args = self.opt.display_models_args.split(',')
            self.opt.usemodel = True
            
        if self.opt.usemodel :
            for model in self._models:
                _opt = copy.deepcopy(self.opt)
                _opt.controller = 'test'
                _opt.model = model
                tmp = get_display_init(model)
                _opt = tmp(_opt)
                for arg in args:
                    tmp = arg.split(':')
                    if model == tmp[0]:
                        _opt.__setattr__(tmp[1],eval(tmp[3])(tmp[2]))
                self.models.append(create_model(_opt))      # create a model given opt.model and other options
                self.models[len(self.models)-1].setup(_opt)               # regular setup: load and print networks; create schedulers

        opt.grab_pcl_fuc = self.get_data
        opt.switch_model_func = self.switch_model
        opt.switch_pcl_func = self.switch_pcl

        self.curpcl = 0
        self.curmodel = 0
        self.datasets.dataset.change_options(self.opt)
        self.vis = CVis(opt)
        self.npoints = opt.pclnpoints
    
    def get_data(self):
        if self.opt.display_3d == 'yes':
            pcl,rgb = self.datasets.dataset.__getitem__(self.curpcl)
            npoints = min(pcl.shape[0],self.npoints)
            if pcl.shape[0] == 0:
                return
            choice = np.random.choice(pcl.shape[0],npoints,False)
            self.vis.changed = True
            
            self.vis.xyz = pcl[choice,:3]
            if rgb is not None:
                self.vis.mapcolor(pcl.shape[0],pcl[choice,3:6])
            else:
                self.vis.mapcolor(self.vis.xyz.shape[0])
            self.vis.change_pcl_range()
            pass
        elif self.opt.display_3d == 'no':
            rgb = self.datasets.dataset.__getitem__(self.curpcl)
            rgb_dict = {}
            rgb_dict['img'] = rgb
            if self.opt.usemodel:
                for model in self.models:
                    rgb_dict=model.display_init(rgb_dict)
                    model.set_input(rgb_dict)
                    model.forward()
                    rgb_dict = model.display_proc(rgb_dict)
            
            self.vis.show = rgb_dict['img']
            pass

    def excute(self):
        # self.vis.run_pcl_dsiplay_process(self.opt.name)
        
        self.vis.run_pcl_dsiplay_process()
        self.datasets.dataset.cam.close()
    
    def switch_model(self,direct):
        if self.models[self.curmodel] is not None:
            if isinstance(self.models[self.curmodel],dict):
                for key in self.models[self.curmodel].keys():
                    self.models[self.curmodel][key].to(self.device)
            else:
                self.models[self.curmodel].to('cpu')
            
        if len(self.models) > 1:
            self.curmodel += direct
            if self.curmodel >= len(self.models):
                self.curmodel=0
            elif self.curmodel < 0:
                self.curmodel += len(self.models)
            self.set_npoints()
                
        if self.models[self.curmodel] is not None:
            if isinstance(self.models[self.curmodel],dict):
                for key in self.models[self.curmodel].keys():
                    self.models[self.curmodel][key].to(self.device)
            else:
                self.models[self.curmodel].to(self.device)
    
    
    def switch_pcl(self,direct):
        datalen = self.datasets.dataset.__len__()
        if datalen > 1:
            self.curpcl += direct
            if self.curpcl >= datalen:
                self.curpcl=0
            elif self.curpcl < 0:
                self.curpcl += datalen
            